from docx import Document
from docx.shared import Pt, RGBColor, Cm, Inches
from docx.oxml.ns import qn
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.enum.table import WD_TABLE_ALIGNMENT
from utils.analysis.mediation.mediation_model14 import MediationModel14, Model14Result
from utils.report.base import format_table, create_table
import pandas as pd


def _generate_report(data, document, table_count):
    table1 = _generate_table1(data, document, table_count)  # 表1
    table3 = _generate_table3(data, document, table_count)  # 表3
    y2x_xz = "" if table1.cell(4, 1).text.endswith("*") else "未"
    y2x_con = "" if table1.cell(4, 1).text.endswith("*") else "不"
    y2xm2_xz = "" if table1.cell(6, 5).text.endswith("*") else "未"
    y2xm2_con = "" if table1.cell(6, 5).text.endswith("*") else "不"
    p1 = (
        f"从上表可知，调节作用分为三个模型，模型1中包括自变量({data.x})，以及{data.covar}等4个控制变量；"
        f"模型2在模型1的基础上加入调节变量({data.w})，模型3在模型2的基础上加入交互项(自变量与调节变量的乘积项)。"
    )
    document.add_paragraph(p1)  # 插入段落
    p2 = (
        f"针对模型1，其目的在于研究在不考虑调节变量({data.w})的干扰时，自变量({data.x})对于因变量({data.y})的影响情况。"
        f"从上表格可知，自变量({data.x})（β={table1.cell(4, 3).text}，t={table1.cell(4, 4).text}）{y2x_xz}呈现出显著性。"
        f"意味着{data.x}对于{data.y}{y2x_con}会产生显著正向影响关系。"
        f"从上表格可知，{data.x}与{data.w}的交互项{y2xm2_xz}呈现出显著性（β={table1.cell(6, 1).text}，t={table1.cell(6, 2).text}）。"
        f"意味着{data.x}对于{data.y}影响时，调节变量({data.w})在不同水平时，影响幅度{y2xm2_con}具有显著性差异。"
    )
    document.add_paragraph(p2)  # 插入段落
    return


def _generate_table1(data, document, table_count):
    # 新建表格
    init_rows = 3
    init_cols = 5
    # table = document.add_table(rows=init_rows, cols=init_cols)
    table = create_table(document, init_rows, init_cols)
    # 1. 填充1-3行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '回归模型汇总表格'
    table.cell(1, 0).merge(table.cell(2, 0))
    table.cell(1, 0).text = "变量"
    # m <- x
    table.cell(1, 1).merge(table.cell(1, 2))
    table.cell(1, 1).text = "模型1;" + data.m
    table.cell(2, 1).text = "β"
    table.cell(2, 2).text = "t"
    # y <- x
    table.cell(1, 3).merge(table.cell(1, 4))
    table.cell(1, 3).text = "模型2;" + data.y
    table.cell(2, 3).text = "β"
    table.cell(2, 4).text = "t"
    # 2. 填充自变量，调节变量
    _add_values(table, "Intercept", data.m2xw.details, data.y2xwm.details)
    _add_values(table, data.x, data.m2xw.details, data.y2xwm.details)
    _add_values(table, data.w, data.m2xw.details, data.y2xwm.details)
    _add_values(table, "INT", data.m2xw.details, data.y2xwm.details)
    _add_values(table, data.m, data.m2xw.details, data.y2xwm.details)
    # 3.填充控制变量
    covars = data.covar
    for i in covars:
        _add_values(table, i, data.m2xw.details, data.y2xwm.details)
    # 4.填充R方，调整后的r方和F值
    add_r2_f(data, table)
    # 格式化表格-所有cell居中
    format_table(table)
    # 增加最后一行，不需要居中
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = '* p<0.05 ** p<0.01 *** p<0.001'
    cur_row.cells[0].merge(cur_row.cells[1]).merge(cur_row.cells[2]).merge(cur_row.cells[3]).merge(cur_row.cells[4])
    return table


def _generate_table2(data, document, table_count):
    document.add_paragraph('')  # 插入段落
    # 新建表格
    init_rows = 5
    init_cols = 8
    table = document.add_table(rows=init_rows, cols=init_cols)
    # 1. 填充1-2行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '调节直接效应结果'
    table.cell(1, 0).text = "水平"

    table.cell(2, 1).text = "低水平（-1SD）"
    table.cell(3, 1).text = "平均值"
    table.cell(4, 1).text = "高水平（+1SD）"

    table.cell(1, 1).text = "水平值"
    table.cell(1, 2).text = "Effect"
    table.cell(1, 3).text = "SE"
    table.cell(1, 4).text = "t值"
    table.cell(1, 5).text = "p值"
    table.cell(1, 6).text = "LLCI"
    table.cell(1, 7).text = "ULCI"
    mediation_table = data.p.direct_model
    table.cell(2, 0).merge(table.cell(3, 0)).merge(table.cell(4, 0)).text = data.m
    if mediation_table._moderators_values:
        cur_2_list = mediation_table._moderators_values[0]
        if cur_2_list and len(cur_2_list) > 2:
            table.cell(2, 2).text = convert_str_to_float_3(cur_2_list[0])
            table.cell(3, 2).text = convert_str_to_float_3(cur_2_list[1])
            table.cell(4, 2).text = convert_str_to_float_3(cur_2_list[2])
    table_res = mediation_table._estimation_results
    _full_table2_values(table, table_res.get('betas'), 2)
    _full_table2_values(table, table_res.get('se'), 3)
    _full_table2_values(table, table_res.get('t'), 4)
    _full_table2_values(table, table_res.get('p'), 5)
    _full_table2_values(table, table_res.get('llci'), 6)
    _full_table2_values(table, table_res.get('ulci'), 7)
    # 格式化表格
    format_table(table)
    return table


def _generate_table3(data, document, table_count):
    document.add_paragraph('')  # 插入段落
    # 新建表格
    init_rows = 5
    init_cols = 7
    # table = document.add_table(rows=init_rows, cols=init_cols)
    table = create_table(document, init_rows, init_cols)
    # 1. 填充1-2行
    # 填充第一行表头
    row = table.rows[0]  # 表格
    f_cell = None
    for i in range(init_cols):  # 合并单元格
        if i == 0:
            f_cell = table.cell(0, 0)
        else:
            f_cell.merge(table.cell(0, i))
    frequency_table_name = ""
    row.cells[0].text = frequency_table_name + '调节间接效应结果'
    table.cell(1, 0).text = "中介变量"

    table.cell(2, 1).text = "低水平（-1SD）"
    table.cell(3, 1).text = "平均值"
    table.cell(4, 1).text = "高水平（+1SD）"

    table.cell(1, 1).text = "水平"
    table.cell(1, 2).text = "水平值"
    table.cell(1, 3).text = "Effect"
    table.cell(1, 4).text = "Boot SE"
    table.cell(1, 5).text = "BootLLCI"
    table.cell(1, 6).text = "BootULCI"
    mediation_table = data.p.indirect_model
    table_res = mediation_table.estimation_results
    table.cell(2, 0).merge(table.cell(3, 0)).merge(table.cell(4, 0)).text = data.m
    if mediation_table._moderators_values:
        cur_2_list = mediation_table._moderators_values[0]
        if cur_2_list and len(cur_2_list) > 2:
            table.cell(2, 2).text = convert_str_to_float_3(cur_2_list[0])
            table.cell(3, 2).text = convert_str_to_float_3(cur_2_list[1])
            table.cell(4, 2).text = convert_str_to_float_3(cur_2_list[2])
    _full_table2_values(table, table_res.get('effect'), 3)
    _full_table2_values(table, table_res.get('se'), 4)
    _full_table2_values(table, table_res.get('llci'), 5)
    _full_table2_values(table, table_res.get('ulci'), 6)
    # 格式化表格
    format_table(table)
    return table


def _full_table2_values(table, cur_list, col):
    if len(cur_list) <= 2 or not col:
        return
    table.cell(4, col).text = convert_str_to_float_3(cur_list[0])
    table.cell(3, col).text = convert_str_to_float_3(cur_list[1])
    table.cell(2, col).text = convert_str_to_float_3(cur_list[2])


def _add_values(table, cur_name, details_m2x, details_y2x):
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = cur_name
    for i in details_m2x:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[1].text = cur_str
            cur_row.cells[2].text = convert_str_to_float_3(i.T)
    for i in details_y2x:
        if cur_name == i.names:
            cur_str = convert_str_to_float_3(i.coef) + add_p_value(i.pval)
            cur_row.cells[3].text = cur_str
            cur_row.cells[4].text = convert_str_to_float_3(i.T)


def add_adj_r2(data, table, name, v1, v2, v3):
    m2x_adj_r2 = 0
    for i in data.r_m2x.details:
        name = i.names
        if 'Intercept' == name:
            m2x_adj_r2 = i.adj_r2
            break
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = name
    cur_row.cells[1].merge(cur_row.cells[2]).text = v1
    cur_row.cells[3].merge(cur_row.cells[4]).text = v2
    cur_row.cells[5].merge(cur_row.cells[6]).text = v3
    return


def add_r2_f(data, table):
    m1 = data.m2xw
    m2 = data.y2xwm
    v1, v2 = 0, 0
    adj_v1, adj_v2 = 0, 0
    for i in m1.details:
        name = i.names
        if 'Intercept' == name:
            v1 = i.r2
            adj_v1 = i.adj_r2
            break
    for i in m2.details:
        name = i.names
        if 'Intercept' == name:
            v2 = i.r2
            adj_v2 = i.adj_r2
    add_last_2_cols(table, "R方", v1, v2)
    add_last_2_cols(table, "调整R2", adj_v1, adj_v2)
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = "F值"
    cur_row.cells[1].merge(cur_row.cells[2]).text = convert_str_to_float_3(m1.f) + add_p_value(m1.p)
    cur_row.cells[3].merge(cur_row.cells[4]).text = convert_str_to_float_3(m2.f) + add_p_value(m2.p)


def add_last_2_cols(table, name, v1, v2):
    table.add_row()  # 表格动态增加一行
    cur_row = table.rows[-1]
    cur_row.cells[0].text = name
    cur_row.cells[1].merge(cur_row.cells[2]).text = convert_str_to_float_3(v1)
    cur_row.cells[3].merge(cur_row.cells[4]).text = convert_str_to_float_3(v2)


def convert_str_to_float_3(value) -> str:
    if not value:
        return ''
    if isinstance(value, str):
        return "{:.3f}".format(float(value))
    elif isinstance(value, float):
        return "{:.3f}".format(value)
    return "{:.3f}".format(value)


def add_p_value(p):
    res = ''
    if p <= 0.001:
        res = '***'
    elif p <= 0.01:
        res = '**'
    elif p <= 0.05:
        res = '*'
    return res


def generate(src_data: [Model14Result], doc: Document = None) -> Document():
    if not doc:
        doc = Document()
    title1 = doc.add_heading(level=1)  # 增加标题
    t1_run = title1.add_run('调节效应分析-模型14')
    t1_run.font.color.rgb = RGBColor(10, 10, 10)
    for i in src_data:
        _generate_report(i, doc, 1)  # 生成表格
    return doc


if __name__ == '__main__':
    # 创建 Document 对象，等价于在电脑上打开一个 Word 文档
    doc = Document()
    pd.set_option('expand_frame_repr', False)
    df = pd.read_excel('./test_datas.xlsx')
    df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # 1.此文件的结果不包含整体的F值和p值
    # 2.若有多个x,y,m，则需要遍历进行获取结果
    x_arr = ['JG']
    y_arr = ['YY']
    m_arr = ['XW']
    w_arr = ['JX']
    covar = ["性别", "年龄", "学历", "企业角色"]
    obj = MediationModel14(df, x_arr, y_arr, m_arr, w_arr, covar)
    src_data = obj.analysis()
    title1 = doc.add_heading(level=1)  # 增加标题
    t1_run = title1.add_run('调节效应分析-模型14')
    t1_run.font.color.rgb = RGBColor(10, 10, 10)
    for i in src_data:
        _generate_report(i, doc, 1)  # 生成表格
    # 保存文档
    doc.save('demo.docx')
